import numpy as np
import matplotlib.pyplot as plt
import scipy.io
import mne  # 替换为MNE库

#******************************************
# 已经测试成功，可以此为1.0版本
#******************************************

def load_eeg_data(file_path):
    """
    从 MAT 文件中加载脑电波数据
    :param file_path: MAT 文件的路径
    :return: 脑电波数据
    """
    try:
        mat_data = scipy.io.loadmat(file_path)
        # 假设脑电波数据存储在 'Impedances_0' 字段中，可根据实际情况修改
        # eeg_data = mat_data['Impedances_0']
        eeg_data = mat_data['data']
        return eeg_data
    except FileNotFoundError:
        print(f"文件 {file_path} 未找到，请检查文件路径。")
        return None
    except KeyError:
        print(f"MAT 文件中未找到 'Impedances_0' 字段，请检查数据结构。")
        return None
    except Exception as e:
        print(f"读取文件时出现错误: {e}")
        return None

def load_bdf_data(file_path):
    """
    从BDF文件中加载脑电波数据（使用MNE库）
    :param file_path: BDF文件路径
    :return: 脑电波数据（numpy数组）
    """
    try:
        raw = mne.io.read_raw_bdf(file_path, preload=True)
        return raw.get_data().T  # 转置以保持与之前相同的维度
    except FileNotFoundError:
        print(f"文件 {file_path} 未找到，请检查文件路径。")
        return None
    except Exception as e:
        print(f"使用MNE读取BDF文件时出现错误: {e}")
        return None

def plot_eeg_channels(eeg_data, sampling_rate, channel_indices):
    """
    绘制指定通道的脑电波信号
    :param eeg_data: 脑电波数据
    :param sampling_rate: 采样率
    :param channel_indices: 要绘制的通道索引列表
    """
    if eeg_data is None:
        return

    num_channels = len(channel_indices)
    fig, axes = plt.subplots(num_channels, 1, figsize=(10, 3 * num_channels))

    if num_channels == 1:
        axes = [axes]

    for idx, ch_idx in enumerate(channel_indices):  # 修复循环变量
        time_vector = np.arange(eeg_data.shape[0]) / sampling_rate
        axes[idx].plot(time_vector, eeg_data[:, ch_idx])  # 修正索引
        axes[idx].set_title(f'EEG Signal of Channel {ch_idx + 1}')  # 显示实际通道号

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    #file_path = 'D:\\dataset\\eeg\\test1\\20250310144145_test_01\\data.mat'
    file_path = 'E:\\20250310144145_test_01\\徐成卓\\徐成卓.bdf'
    sampling_rate = 256
    # 选择要绘制的通道索引，这里选择前 5 个通道，可根据需要修改
    channel_indices = [0, 1, 2, 3, 4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31]
    channel_indices = [0, 1, 2]
    # channel_indices = [0]

    #eeg_data = load_eeg_data(file_path)
    eeg_data = load_bdf_data(file_path)
    plot_eeg_channels(eeg_data, sampling_rate, channel_indices)